import torch
from torch import nn
from torchvision.datasets import ImageFolder  # 直接加载数据目录，进行分类
from torchvision.transforms import transforms  # 将图像进行一定处理的工具
from torch.utils.data import DataLoader

# ==================== 图像数据集处理 ===================
# 对单张图片进行处理
transform = transforms.Compose([
    transforms.ToTensor()  # 将加载的图像数据转化为pytorch.tensor类型
])

# ImageFoler将指定目录下的文件数据进行自动分类
dataset = ImageFolder("./data/numbers", transform=transform)

# 将上述数据集进行 将数据分批次、将数据进行随机打散
dataloader = DataLoader(dataset, batch_size=1000, shuffle=True)
